package org.openmetadata.service.rules;

import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty;
import static org.openmetadata.service.Entity.TEAM;
import static org.openmetadata.service.rules.JsonLogicUtils.evaluateUserInRole;

import com.fasterxml.jackson.core.type.TypeReference;
import io.github.jamsesso.jsonlogic.JsonLogic;
import io.github.jamsesso.jsonlogic.ast.JsonLogicArray;
import io.github.jamsesso.jsonlogic.evaluator.JsonLogicEvaluationException;
import io.github.jamsesso.jsonlogic.evaluator.JsonLogicEvaluator;
import io.github.jamsesso.jsonlogic.evaluator.JsonLogicExpression;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;
import org.openmetadata.common.utils.CommonUtil;
import org.openmetadata.schema.entity.domains.DataProduct;
import org.openmetadata.schema.type.EntityReference;
import org.openmetadata.schema.type.Include;
import org.openmetadata.schema.type.TagLabel;
import org.openmetadata.schema.utils.JsonUtils;
import org.openmetadata.service.Entity;

public class LogicOps {

  public enum CustomLogicOps {
    LENGTH("length"),
    CONTAINS("contains"),
    IS_REVIEWER("isReviewer"),
    IS_OWNER("isOwner"),
    IS_UPDATED_BEFORE("isUpdatedBefore"),
    IS_UPDATED_AFTER("isUpdatedAfter");

    public final String key;

    CustomLogicOps(String param) {
      this.key = param;
    }
  }

  public static void addCustomOps(JsonLogic jsonLogic) {
    // This method can be used to add custom operations that are exposed in the UI.
    // Currently, it adds the LENGTH operation.
    addCustomPublicOps(jsonLogic);
    addCustomPrivateOps(jsonLogic);
  }

  private static void addCustomPublicOps(JsonLogic jsonLogic) {
    jsonLogic.addOperation(
        CustomLogicOps.LENGTH.key,
        (args) -> {
          if (nullOrEmpty(args)) {
            return 0;
          }
          return args.length;
        });

    jsonLogic.addOperation(
        "contains",
        (args) -> {
          if (nullOrEmpty(args) || args.length < 2) {
            return false;
          }

          Object value = args[0];
          Object container = args[1];

          // If either value or container is null/empty, the rule is broken
          if (CommonUtil.nullOrEmpty(value) || CommonUtil.nullOrEmpty(container)) {
            return false;
          }

          if (container instanceof List<?> list) {
            return list.contains(value);
          } else if (container.getClass().isArray()) {
            Object[] array = (Object[]) container;
            return Arrays.asList(array).contains(value);
          } else {
            return container.toString().contains(value.toString());
          }
        });

    // {"isReviewer": { var: "updatedBy"} }
    jsonLogic.addOperation(
        new JsonLogicExpression() {
          @Override
          public String key() {
            return CustomLogicOps.IS_REVIEWER.key;
          }

          @Override
          public Object evaluate(
              JsonLogicEvaluator evaluator, JsonLogicArray arguments, Object data)
              throws JsonLogicEvaluationException {
            return evaluateUserInRole(evaluator, arguments, data, "reviewers");
          }
        });

    // {"isOwner": {var: "updatedBy"} }
    jsonLogic.addOperation(
        new JsonLogicExpression() {
          @Override
          public String key() {
            return CustomLogicOps.IS_OWNER.key;
          }

          @Override
          public Object evaluate(
              JsonLogicEvaluator evaluator, JsonLogicArray arguments, Object data)
              throws JsonLogicEvaluationException {
            return evaluateUserInRole(evaluator, arguments, data, "owners");
          }
        });

    // {"isUpdatedBefore": 1609459200000} - Check if entity was updated before timestamp
    jsonLogic.addOperation(
        new JsonLogicExpression() {
          @Override
          public String key() {
            return CustomLogicOps.IS_UPDATED_BEFORE.key;
          }

          @Override
          public Object evaluate(
              JsonLogicEvaluator evaluator, JsonLogicArray arguments, Object data)
              throws JsonLogicEvaluationException {
            return JsonLogicUtils.evaluateIsUpdatedBefore(evaluator, arguments, data);
          }
        });

    // {"isUpdatedAfter": 1609459200000} - Check if entity was updated after timestamp
    jsonLogic.addOperation(
        new JsonLogicExpression() {
          @Override
          public String key() {
            return CustomLogicOps.IS_UPDATED_AFTER.key;
          }

          @Override
          public Object evaluate(
              JsonLogicEvaluator evaluator, JsonLogicArray arguments, Object data)
              throws JsonLogicEvaluationException {
            return JsonLogicUtils.evaluateIsUpdatedAfter(evaluator, arguments, data);
          }
        });
  }

  /**
   * Here we'll add validations that are not going to be exposed in the UI
   * in the rule settings.
   * This is why we are not picking the names from the CustomLogicOps enum.
   */
  private static void addCustomPrivateOps(JsonLogic jsonLogic) {
    // This is expecting the `owners` to be passed to the rule argument
    jsonLogic.addOperation(
        "multipleUsersOrSingleTeamOwnership",
        (args) -> {
          if (nullOrEmpty(args) || Arrays.stream(args).allMatch(CommonUtil::nullOrEmpty)) {
            return true;
          }
          List<EntityReference> owners =
              JsonUtils.convertValue(args, new TypeReference<List<EntityReference>>() {});
          if (nullOrEmpty(owners)) {
            return true;
          }
          long teamCount = owners.stream().filter(owner -> owner.getType().equals(TEAM)).count();
          long userCount = owners.size() - teamCount;
          if (teamCount > 1 || (teamCount > 0 && userCount > 0)) {
            return false; // More than one team or both team and user ownership
          }
          return true;
        });

    // Domain validation for Data Product assignment
    // This validates that entities being assigned to Data Products have matching domains
    jsonLogic.addOperation(
        "validateDataProductDomainMatch",
        (args) -> {
          if (nullOrEmpty(args) || args.length < 2) {
            return true; // If no data products or domains, validation passes
          }

          List<EntityReference> dataProducts =
              JsonUtils.convertValue(args[0], new TypeReference<List<EntityReference>>() {});
          List<EntityReference> domains =
              JsonUtils.convertValue(args[1], new TypeReference<List<EntityReference>>() {});

          if (nullOrEmpty(dataProducts)) {
            return true; // If no data products, validation passes
          } else if (nullOrEmpty(domains)) {
            return false; // If data products but no domains, validation fails
          }

          // Convert entity domains to a Set of UUIDs for efficient lookup
          Set<UUID> entityDomainIds =
              domains.stream().map(EntityReference::getId).collect(Collectors.toSet());

          // Get all data product entities in bulk instead of using a loop
          try {
            List<DataProduct> dpEntities =
                Entity.getEntities(dataProducts, "domains", Include.NON_DELETED);

            for (DataProduct dpEntity : dpEntities) {
              List<EntityReference> dpDomains = dpEntity.getDomains();

              if (nullOrEmpty(dpDomains)) {
                continue; // If data product has no domains, skip validation
              }

              // Convert data product domains to a Set of UUIDs for efficient comparison
              Set<UUID> dpDomainIds =
                  dpDomains.stream().map(EntityReference::getId).collect(Collectors.toSet());

              // Use Collections.disjoint for O(n+m) performance instead of O(n*m)
              boolean hasMatchingDomain = !Collections.disjoint(entityDomainIds, dpDomainIds);

              if (!hasMatchingDomain) {
                return false; // Domain mismatch found
              }
            }
          } catch (Exception e) {
            // If we can't fetch the data products, fail validation for safety
            return false;
          }
          return true; // All data products have matching domains
        });

    // Example: {"filterReferenceByType":[{"var":"owner"},"team"]}
    jsonLogic.addOperation(
        "filterReferenceByType",
        (args) -> {
          if (nullOrEmpty(args)) {
            return false;
          }
          String type = (String) args[1];
          if (nullOrEmpty(args[0]) || CommonUtil.nullOrEmpty(type)) {
            return List.of();
          }

          List<EntityReference> refs =
              JsonUtils.convertValue(args[0], new TypeReference<List<EntityReference>>() {});
          return refs.stream().filter(ref -> ref.getType().equals(type)).toList();
        });

    // Example: {"filterTagsBySource":[{"var":"tags"},"Glossary"]}
    jsonLogic.addOperation(
        "filterTagsBySource",
        (args) -> {
          if (nullOrEmpty(args)) {
            return List.of();
          }
          String type = (String) args[1];

          if (nullOrEmpty(args[0]) || CommonUtil.nullOrEmpty(type)) {
            return List.of();
          }

          List<TagLabel> tags =
              JsonUtils.convertValue(args[0], new TypeReference<List<TagLabel>>() {});
          return tags.stream().filter(ref -> ref.getSource().toString().equals(type)).toList();
        });
  }

  /**
   * Returns a list of keys for custom logic operations.
   * This is exposed in /system/settings/customLogicOps for the UI.
   */
  public static List<String> getCustomOpsKeys() {
    return Arrays.stream(CustomLogicOps.values()).map(op -> op.key).toList();
  }
}
